import cv2
import numpy as np
import matplotlib.pyplot as plt

# 均值滤波
def mean_filter(src, ksize):
    dst = np.zeros(src.shape, dtype=np.uint8)
    h, w = src.shape[0], src.shape[1]
    rx = ksize[0] // 2
    ry = ksize[1] // 2
    nc = src.shape[-1]
    for y in range(h):
        for x in range(w):
            n = 0
            v = np.zeros(nc, dtype=np.float32)
            for i in range(-ry, ry+1):
                iy = i + y
                if iy < 0 or iy >= h:
                    continue
                for j in range(-rx, rx+1):
                    jx = j + x
                    if jx < 0 or jx >= w:
                        continue
                    v += src[iy][jx]
                    n += 1
            v = v / n
            dst[y][x] = v.astype(np.uint8)

    return dst


# 选择均值滤波
def select_mean_filter(src, ksize, threshold):
    dst = np.zeros(src.shape, dtype=np.uint8)
    h, w = src.shape[0], src.shape[1]
    rx = ksize[0] // 2
    ry = ksize[1] // 2
    nc = src.shape[-1]
    for y in range(h):
        for x in range(w):
            n = 0
            v = np.zeros(nc, dtype=np.float32)
            for i in range(-ry, ry+1):
                iy = i + y
                if iy < 0 or iy >= h:
                    continue
                for j in range(-rx, rx+1):
                    jx = j + x
                    if jx < 0 or jx >= w:
                        continue
                    v += src[iy][jx]
                    n += 1
            v = (v / n).astype(np.uint8)
            mask = v - src[y][x] < threshold
            v[mask] = src[y][x][mask]
            dst[y][x] = v

    return dst

# 中值滤波
def median_filter(src, ksize):
    dst = np.zeros(src.shape, dtype=np.uint8)
    h, w = src.shape[0], src.shape[1]
    rx = ksize[0] // 2
    ry = ksize[1] // 2
    nc = src.shape[-1]
    for y in range(h):
        for x in range(w):
            n = 0
            v = np.empty(((2*rx+1)*(2*ry+1), nc), dtype=np.uint8)
            for i in range(-ry, ry + 1):
                iy = i + y
                if iy < 0 or iy >= h:
                    continue
                for j in range(-rx, rx + 1):
                    jx = j + x
                    if jx < 0 or jx >= w:
                        continue
                    v[n] = src[iy][jx]
                    n += 1
            v = np.resize(v, (n, nc))
            v.sort(axis=0)
            dst[y][x] = v[n//2]

    return dst

filename = 'lena.bmp'
src = cv2.imread(filename, cv2.IMREAD_COLOR)
src = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
# 滤波处理
#dst = mean_filter(src, (5, 5))
#dst = select_mean_filter(src, (5, 5), 80)
#dst = median_filter(src, (5, 5))
# OpenCV均值滤波
#dst = cv2.boxFilter(src, src.shape[2], (5, 5))
# OpenCV高斯滤波
#dst = cv2.GaussianBlur(src, (5, 5), sigmaX=2., sigmaY=2.)
# k1 = cv2.getGaussianKernel(5, sigma=2)
# kernel = k1 * k1.T
# dst = cv2.filter2D(src, src.shape[2], kernel)
# OpenCV中值滤波
#dst = cv2.medianBlur(src, 5)
# OpenCV双边滤波
#dst = cv2.bilateralFilter(dst, 20, sigmaColor=35, sigmaSpace=10)
# OpenCV均值漂移滤波
#dst = cv2.pyrMeanShiftFiltering(dst, sp=5, sr=20);
# 图像锐化
kernel = np.array([-1, -2, -1, -2, 19, -2, -1, -2, -1]) / 7.0
dst = cv2.filter2D(src, src.shape[2], kernel)
# 显示原始图像
ax1 = plt.subplot(121)
ax1.set_title('Src')
plt.imshow(src)
plt.xticks([]), plt.yticks([])
# 显示滤波后图像
ax2 = plt.subplot(122)
ax2.set_title('Filter')
plt.imshow(dst)
plt.xticks([]), plt.yticks([])
plt.show()
